package com.uh.rds.testing.validator;

import com.uh.rds.testing.base.DataValidator;
import com.uh.rds.testing.base.DataValidatorStep;
import com.uh.rds.testing.conn.Endpoint;
import com.uh.rds.testing.conn.RdsConnectInfo;
import com.uh.rds.testing.conn.Shard;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.util.*;

import static com.uh.rds.testing.TestingConstants.THREADS_WAIT_JOIN;
import static com.uh.rds.testing.data.TestDataLoader.readSubsetData;
import static com.uh.rds.testing.data.TestDataLoader.splitMapList;
import com.uh.rds.testing.conn.ConnectionMode;

/**
 * 多线程并发的数据验证运行器
 */
public class DataValidatorRunner {
    private static Logger logger = LoggerFactory.getLogger(DataValidatorRunner.class);

    private List<ValidatorTypeBuilder> validatorBuilders = new ArrayList<>();

    private List<DataValidatorThread> threads;

    private DataValidatorArguments config;

    private boolean ready = false;

    private List<Thread> startedThreads; //启动中的线程

    private DataValidatorStep[] startedSteps;  //启动中的执行步骤

    public DataValidatorRunner() {}

    public DataValidatorRunner(DataValidatorArguments config) {
        this.config = config;
    }

    public void addValidator(Class<? extends DataValidator> clazz, File dataFile) {
        validatorBuilders.add(new ValidatorTypeBuilder(clazz, dataFile));
    }

    public void setConfig(DataValidatorArguments config) {
        this.config = config;
    }

    public DataValidatorArguments getConfig() {
        return config;
    }

    /**
     * 准备测试线程和数据
     * @param connection 连接信息
     */
    public void prepareThreadsData(RdsConnectInfo connection) {
        this.ready = false;
        this.threads = new ArrayList<>();
        ConnectionMode mode = connection.getMode();
        if(mode == ConnectionMode.CLUSTER) {
            prepareClusterThreadsData(connection);
        }
        else if(mode == ConnectionMode.MASTER || mode == ConnectionMode.MASTER_SLAVE) {
            prepareSingleThreadsData(connection);
        }
        else {
            throw new RuntimeException("Not support deploy mode:" + mode);
        }

        this.ready = true;
    }

    /**
     * 是否准备好
     * @return
     */
    public boolean isReady() {
        return ready;
    }

    public void startTest(DataValidatorStep... steps) {
        if(!isReady()) {
            throw new RuntimeException("Not ready! Please call prepareThreadsData(connection) first!");
        }

        if(startedThreads != null && !startedThreads.isEmpty()) {
            throw new RuntimeException("Threads already started, please call runTest() instead of startTest()!");
        }

        this.startedSteps = steps;
        startedThreads = new ArrayList<>();
        //启动所有线程
        for(DataValidatorThread t : threads) {
            t.setSteps(steps);

            //如果有大批thread创建是，会使线程和socket连接数暴增，导致连接数过多，所以这里sleep一下
            try {
                Thread.sleep(config.threadCreateInterval);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }

            Thread thread = new Thread(t);
            thread.start();
            startedThreads.add(thread);
        }
    }

    public DataValidatorResult waitTestResult() {

        try {
            Thread.sleep(THREADS_WAIT_JOIN);
        } catch (InterruptedException e) {
            logger.error("Thread sleep error!", e);
        }

        //等待所有线程停止
        for(Thread t : startedThreads) {

            try {
                t.join();
            } catch (InterruptedException e) {
                logger.error("Thread join error!", e);
            }
        }

        DataValidatorResult ret = new DataValidatorResult(startedSteps, threads);

        //清空状态
        this.startedSteps = null;
        this.startedThreads = null;

        return ret;
    }

    /**
     * 运行测试
     * @param steps 测试步骤
     * @return 测试结果
     */
    public DataValidatorResult runTest(DataValidatorStep... steps) {
        startTest(steps);
        return waitTestResult();
    }

    private void prepareSingleThreadsData(RdsConnectInfo connection) {

        List<Endpoint> masters = connection.getMasterEndpoints();
        List<Endpoint> slaves = connection.getSlaveEndpoints();

        for(int i = 0 ; i < config.threadCount ; i++) {
            Endpoint endpoint = masters.get(i % masters.size());
            Endpoint slaveEndpoint = slaves.isEmpty() ? null : slaves.get(i % slaves.size());
            DataValidatorThread thread = new DataValidatorThread(i, endpoint, slaveEndpoint, config);
            threads.add(thread);
        }

        //把每种类型的数据从磁盘中读取，平局分配到不同的thread的数据中。
        for(ValidatorTypeBuilder validType : validatorBuilders) {
            //从磁盘读取对应类型的数据
            Map<String, List<String[]>> typeData = validType.getData();
            List<Map<String, List<String[]>>> threadsData = splitMapList(typeData, threads.size());
            // 为每一个thread的type下设置数据。
            int i = 0;
            for(DataValidatorThread thread : threads) {
                DataValidator validator = validType.newValidator();
                validator.init(thread.getIndex(), -1, config, threadsData.get(i));
                thread.addValidator(validator);
                i ++;
            }
        }

    }


    private void prepareClusterThreadsData(RdsConnectInfo connection) {
        int threadSize = config.threadCount;

        List<Shard> shards = connection.getShards();
        int shardCount = shards.size();

        // 创建shards数量的thread列表
        List<List<DataValidatorThread>> shardsThreads = new ArrayList<>(shardCount);
        for (int i = 0; i < shardCount; i++) {  shardsThreads.add(new ArrayList<>()); }

        // 1. 创建Thread列表, 放入: threads(所有线程)中,  和shardsThreads(按shardIndex拆分的线程子集)中。
        int subThreadSize = threadSize / shardCount; //每个分片平均分配几个Thread
        int shardIdx = 0;
        logger.info("Shard size:" + shardCount + ", thread size:" + threadSize);
        // 循环创建指定数量的thread，并把它们分配到各个shard下（List<List<DataValidatorThread>> shardsThreads）
        for(int i = 0 ; i < threadSize ; i++) {
            if( i!=0 && (i % subThreadSize == 0) && (shardIdx<shardCount-1)) {
                shardIdx++;
            }

            Shard shardInfo = shards.get(shardIdx);
            Endpoint masterJedis = shardInfo.getMaster();
            Endpoint slaveJedis =  shardInfo.getSlaves().isEmpty() ? null : shardInfo.getSlaves().get(i % shardInfo.getSlaves().size());
            DataValidatorThread thread = new DataValidatorThread(i, masterJedis, slaveJedis, config);
            threads.add(thread);
            List<DataValidatorThread> shardThreads = shardsThreads.get(shardIdx);
            shardThreads.add(thread);
        }

        // 2. 把每种类型的数据从磁盘中读取，分配到不同的shard 》再分配到其下的threads中。
        for(ValidatorTypeBuilder validType : validatorBuilders) {
            Map<String, List<String[]>> typeData = validType.getData(); //从磁盘读取对应类型的数据
            List<Map<String, List<String[]>>> shardsData = new ArrayList<>(shardCount);
            for(int i=0 ; i<shardCount ; i++) { shardsData.add(new HashMap<String, List<String[]>>()); }

            //把数据分配到不同的shard中
            for(Map.Entry<String, List<String[]>> entry : typeData.entrySet()) {
                String key = entry.getKey();
                List<String[]> value = entry.getValue();
                for(int i=0 ; i<shardCount ; i++) {
                    Shard shard = shards.get(i);
                    if(shard.isKeyInShard(key)) {
                        shardsData.get(i).put(key, value);
                        break;
                    }
                }
            }

            //再对每个分片中的数据平均分配到其下的threads中
            for(int i=0 ; i<shardCount ; i++) {
                Shard shardInfo = shards.get(i);
                Map<String, List<String[]>> shardData = shardsData.get(i);
                logger.info("    Type:{} Shard{} data size:{}", validType.getType(), shardInfo.getIndex(), shardData.size());
                List<DataValidatorThread> subThreads = shardsThreads.get(i);
                List<Map<String, List<String[]>>> subThreadsData = splitMapList(shardData, subThreads.size());


                // 为每一个thread的type下设置数据。
                int j = 0;
                for (DataValidatorThread thread : subThreads) {
                    DataValidator validator = validType.newValidator();
                    validator.init(thread.getIndex(), shardInfo.getIndex(), config, subThreadsData.get(j));
                    thread.addValidator(validator);
                    j++;
                }
            }
        }
    }



    private static class ValidatorTypeBuilder {
        public ValidatorTypeBuilder(Class<? extends DataValidator> clazz, File dataFile) {
            try {
                DataValidator instance = clazz.newInstance();
                this.type = instance.getType();
                this.dataFile = dataFile;
                this.clazz = clazz;
            } catch (InstantiationException | IllegalAccessException e) {
                logger.error("Create ValidatorType error!", e);
            }
        }

        public DataValidator newValidator() {
            DataValidator instance = null;
            try {
                instance = clazz.newInstance();
            } catch (InstantiationException | IllegalAccessException e) {
                logger.error("Create ValidatorType error!", e);
            }

            return instance;
        }

        public Map<String, List<String[]>> getData() {
            if(data == null) {
                if(dataFile == null) {
                    throw new IllegalArgumentException("Data file is null for type:" + type);
                }
                if(!dataFile.exists()) {
                    throw new IllegalArgumentException("Data file not exists:" + dataFile.getAbsolutePath());
                }
                data = readSubsetData(dataFile); //从磁盘读取对应类型的数据
                logger.info("Type " + type + " " + data.size() + " rows Data loaded!");
            }
            return data;
        }

        public String getType() {
            return type;
        }

        private Map<String, List<String[]>> data = null;
        private String type;
        private File dataFile;
        private Class<? extends DataValidator> clazz;
    }

}
